/**********************************************************************
  Filename    : ESP32-S3 Camera WebSocket Client
  Description : ESP32-S3 OV2640 camera streaming to WebSocket server
  Based on    : Original Python video capture system
  Modification: 2025/06/22
**********************************************************************/

#include "esp_camera.h"
#include <WiFi.h>
#include <WebSocketsClient.h>
#include <ArduinoJson.h>
#include <base64.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_timer.h"

// ===================
// Camera model selection
// ===================
#define CAMERA_MODEL_ESP32S3_EYE // Has PSRAM
#include "camera_pins.h"

// ===================
// WiFi credentials
// ===================
const char* ssid     = "TP-LINK_House_5G";
const char* password = "123455678";

// ===================
// Server configuration
// ===================
const char* server_host = "XXX.XXX.XXX.XXX";
const int server_port = 5004;

// ===================
// Camera configuration
// ===================
#define CAMERA_ID "ESP32-S3Camera"
#define TARGET_FPS 5
#define JPEG_QUALITY 45
#define FRAME_SIZE FRAMESIZE_96X96  // 320x240

// ===================
// Global variables
// ===================
WebSocketsClient webSocket;
bool is_connected = false;
bool is_capturing = false;
bool is_running = true;
uint32_t frame_count = 0;
unsigned long start_time = 0;
unsigned long last_frame_time = 0;
unsigned long last_reconnect_attempt = 0;
const unsigned long reconnect_interval = 5000; // 5 seconds

// Task handles
TaskHandle_t cameraTaskHandle = NULL;
TaskHandle_t reconnectTaskHandle = NULL;

// ===================
// Function declarations
// ===================
void initCamera();
void connectToWiFi();
void connectToServer();
void webSocketEvent(WStype_t type, uint8_t * payload, size_t length);
void sendCameraRegistration();
void sendControlResponse(const char* command, bool success, const char* message);
void handleControlCommand(const char* command);
void captureAndSendFrame();
void cameraTask(void* parameter);
void reconnectTask(void* parameter);
void printStats();

void setup() {
  Serial.begin(115200);
  Serial.setDebugOutput(true);
  Serial.println();
  Serial.println("=== ESP32-S3 Camera WebSocket Client ===");

  // Initialize camera
  initCamera();
  
  // Connect to WiFi
  connectToWiFi();
  
  // Setup WebSocket
  webSocket.begin(server_host, server_port, "/");
  webSocket.onEvent(webSocketEvent);
  webSocket.setReconnectInterval(5000);
  webSocket.enableHeartbeat(15000, 3000, 2);
  
  // Create camera streaming task
  xTaskCreatePinnedToCore(
    cameraTask,
    "CameraTask",
    8192,
    NULL,
    1,
    &cameraTaskHandle,
    0  // Core 0
  );
  
  // Create reconnection task
  xTaskCreatePinnedToCore(
    reconnectTask,
    "ReconnectTask", 
    4096,
    NULL,
    1,
    &reconnectTaskHandle,
    1  // Core 1
  );
  
  Serial.println("ESP32-S3 Camera Client Started");
  Serial.println("Waiting for server connection...");
}

void loop() {
  webSocket.loop();
  delay(10);
}

void initCamera() {
  Serial.println("Initializing camera...");
  
  camera_config_t config;
  config.ledc_channel = LEDC_CHANNEL_0;
  config.ledc_timer   = LEDC_TIMER_0;
  config.pin_d0       = Y2_GPIO_NUM;
  config.pin_d1       = Y3_GPIO_NUM;
  config.pin_d2       = Y4_GPIO_NUM;
  config.pin_d3       = Y5_GPIO_NUM;
  config.pin_d4       = Y6_GPIO_NUM;
  config.pin_d5       = Y7_GPIO_NUM;
  config.pin_d6       = Y8_GPIO_NUM;
  config.pin_d7       = Y9_GPIO_NUM;
  config.pin_xclk     = XCLK_GPIO_NUM;
  config.pin_pclk     = PCLK_GPIO_NUM;
  config.pin_vsync    = VSYNC_GPIO_NUM;
  config.pin_href     = HREF_GPIO_NUM;
  config.pin_sscb_sda = SIOD_GPIO_NUM;
  config.pin_sscb_scl = SIOC_GPIO_NUM;
  config.pin_pwdn     = PWDN_GPIO_NUM;
  config.pin_reset    = RESET_GPIO_NUM;
  config.xclk_freq_hz = 30000000;
  config.frame_size   = FRAME_SIZE;
  config.pixel_format = PIXFORMAT_JPEG;
  config.grab_mode    = CAMERA_GRAB_WHEN_EMPTY;
  config.fb_location  = CAMERA_FB_IN_PSRAM;
  config.jpeg_quality = JPEG_QUALITY;
  config.fb_count     = 1;
  
  if(psramFound()){
    config.jpeg_quality = JPEG_QUALITY;
    config.fb_count = 2;
    config.grab_mode = CAMERA_GRAB_LATEST;
    Serial.println("PSRAM found, using optimized settings");
  } else {
    config.frame_size = FRAMESIZE_SVGA;
    config.fb_location = CAMERA_FB_IN_DRAM;
    Serial.println("No PSRAM, using DRAM");
  }

  esp_err_t err = esp_camera_init(&config);
  if (err != ESP_OK) {
    Serial.printf("Camera init failed with error 0x%x\n", err);
    ESP.restart();
    return;
  }

  sensor_t * s = esp_camera_sensor_get();
  s->set_vflip(s, 1);
  s->set_brightness(s, 1);
  s->set_saturation(s, 0);
  
  Serial.println("Camera initialized successfully");
}

void connectToWiFi() {
  Serial.printf("Connecting to WiFi: %s\n", ssid);
  WiFi.begin(ssid, password);
  WiFi.setSleep(false);

  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    delay(500);
    Serial.print(".");
    attempts++;
  }
  
  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("");
    Serial.println("WiFi connected");
    Serial.print("IP address: ");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println("");
    Serial.println("WiFi connection failed! Restarting...");
    ESP.restart();
  }
}

void webSocketEvent(WStype_t type, uint8_t * payload, size_t length) {
  switch(type) {
    case WStype_DISCONNECTED:
      Serial.println("WebSocket Disconnected");
      is_connected = false;
      break;
      
    case WStype_CONNECTED:
      Serial.printf("WebSocket Connected to: %s\n", payload);
      is_connected = true;
      sendCameraRegistration();
      break;
      
    case WStype_TEXT:
      {
        Serial.printf("Received: %s\n", payload);
        
        // Parse JSON message
        DynamicJsonDocument doc(1024);
        deserializeJson(doc, payload);
        
        const char* type_str = doc["type"];
        if (strcmp(type_str, "control_command") == 0) {
          const char* command = doc["command"];
          handleControlCommand(command);
        } else if (strcmp(type_str, "register_response") == 0) {
          bool success = doc["success"];
          if (success) {
            Serial.println("Registration confirmed by server");
          } else {
            Serial.printf("Registration failed: %s\n", doc["message"].as<const char*>());
          }
        }
      }
      break;
      
    case WStype_BIN:
      Serial.printf("Received binary data: %u bytes\n", length);
      break;
      
    case WStype_ERROR:
      Serial.printf("WebSocket Error: %s\n", payload);
      break;
      
    default:
      break;
  }
}

void sendCameraRegistration() {
  if (!is_connected) return;
  
  DynamicJsonDocument doc(512);
  doc["type"] = "camera_register";
  doc["camera_id"] = CAMERA_ID;
  doc["timestamp"] = millis() / 1000.0;
  
  String message;
  serializeJson(doc, message);
  
  webSocket.sendTXT(message);
  Serial.println("Camera registration sent");
}

void sendControlResponse(const char* command, bool success, const char* message) {
  if (!is_connected) return;
  
  DynamicJsonDocument doc(512);
  doc["type"] = "control_response";
  doc["command"] = command;
  doc["success"] = success;
  doc["message"] = message;
  doc["camera_id"] = CAMERA_ID;
  doc["timestamp"] = millis() / 1000.0;
  
  String response;
  serializeJson(doc, response);
  
  webSocket.sendTXT(response);
}

void handleControlCommand(const char* command) {
  Serial.printf("Handling command: %s\n", command);
  
  if (strcmp(command, "start_capture") == 0) {
    if (!is_capturing) {
      is_capturing = true;
      frame_count = 0;
      start_time = millis();
      Serial.println("Started video capture");
      sendControlResponse(command, true, "Video capture started");
    } else {
      Serial.println("Capture already running");
      sendControlResponse(command, true, "Video capture already running");
    }
  } 
  else if (strcmp(command, "stop_capture") == 0) {
    if (is_capturing) {
      is_capturing = false;
      Serial.println("Stopped video capture");
      sendControlResponse(command, true, "Video capture stopped");
    } else {
      Serial.println("Capture already stopped");
      sendControlResponse(command, true, "Video capture already stopped");
    }
  }
  else {
    Serial.printf("Unknown command: %s\n", command);
    sendControlResponse(command, false, "Unknown command");
  }
}

void captureAndSendFrame() {
  if (!is_connected || !is_capturing) return;
  
  camera_fb_t * fb = esp_camera_fb_get();
  if (!fb) {
    Serial.println("Camera capture failed");
    return;
  }
  
  // Encode frame to base64
  String base64Frame = base64::encode(fb->buf, fb->len);
  
  // Create JSON message
  DynamicJsonDocument doc(base64Frame.length() + 512);
  doc["type"] = "frame_data";
  doc["camera_id"] = CAMERA_ID;
  doc["frame"] = base64Frame;
  doc["timestamp"] = millis() / 1000.0;
  doc["frame_count"] = frame_count;
  
  String message;
  serializeJson(doc, message);
  
  // Send frame
  if (webSocket.sendTXT(message)) {
    frame_count++;
    last_frame_time = millis();
    
    // Print stats every 50 frames
    if (frame_count % 50 == 0) {
      printStats();
    }
  } else {
    Serial.println("Failed to send frame");
  }
  
  esp_camera_fb_return(fb);
}

void printStats() {
  if (start_time > 0) {
    unsigned long elapsed = millis() - start_time;
    float avg_fps = (float)frame_count / (elapsed / 1000.0);
    Serial.printf("Frames sent: %d, Average FPS: %.2f\n", frame_count, avg_fps);
    
    // Memory info
    Serial.printf("Free heap: %d bytes, Min free heap: %d bytes\n", 
                  ESP.getFreeHeap(), ESP.getMinFreeHeap());
  }
}

void cameraTask(void* parameter) {
  const TickType_t frameDelay = pdMS_TO_TICKS(1000 / TARGET_FPS);
  TickType_t lastWakeTime = xTaskGetTickCount();
  
  Serial.println("Camera task started");
  
  while (is_running) {
    if (is_connected && is_capturing) {
      captureAndSendFrame();
    }
    
    vTaskDelayUntil(&lastWakeTime, frameDelay);
  }
  
  Serial.println("Camera task ended");
  vTaskDelete(NULL);
}

void reconnectTask(void* parameter) {
  Serial.println("Reconnect task started");
  
  while (is_running) {
    // Check WiFi connection
    if (WiFi.status() != WL_CONNECTED) {
      Serial.println("WiFi disconnected, attempting reconnection...");
      WiFi.reconnect();
      vTaskDelay(pdMS_TO_TICKS(5000));
      continue;
    }
    
    // Check WebSocket connection
    if (!is_connected) {
      unsigned long current_time = millis();
      if (current_time - last_reconnect_attempt > reconnect_interval) {
        Serial.println("Attempting WebSocket reconnection...");
        webSocket.disconnect();
        vTaskDelay(pdMS_TO_TICKS(1000));
        webSocket.begin(server_host, server_port, "/");
        last_reconnect_attempt = current_time;
      }
    }
    
    vTaskDelay(pdMS_TO_TICKS(2000));
  }
  
  Serial.println("Reconnect task ended");
  vTaskDelete(NULL);
}
